package com.itqiqi.api.tableapi.udf;

import com.itqiqi.api.pojo.SensorReading;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.streaming.api.TimeCharacteristic;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.table.api.Table;
import org.apache.flink.table.api.java.StreamTableEnvironment;
import org.apache.flink.table.functions.AggregateFunction;
import org.apache.flink.table.functions.TableFunction;
import org.apache.flink.types.Row;

public class UdfTest3_AggregateFunction {

    public static void main(String[] args) throws Exception {

        // 创建环境
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);
        env.setStreamTimeCharacteristic(TimeCharacteristic.EventTime);

        // 创建表环境
        StreamTableEnvironment tableEnv = StreamTableEnvironment.create(env);

        // 读取数据
        DataStreamSource<String> inputStream = env.readTextFile("input/sensor.txt");

        SingleOutputStreamOperator<SensorReading> dataStream = inputStream.map(new MapFunction<String, SensorReading>() {
            @Override
            public SensorReading map(String s) throws Exception {
                String[] words = s.split(",");
                return new SensorReading(words[0], new Long(words[1]), new Double(words[2]));
            }
        });

        // 将流转换成表
        Table sensorTable = tableEnv.fromDataStream(dataStream, "id, timestamp as ts, temperature as temp");

        // TODO: 2022/5/19 自定义聚合函数，求当前传感器的平均温度值
        // 在环境中注册UDF
        AvgTemp avgTemp = new AvgTemp();
        tableEnv.registerFunction("avgTemp", avgTemp);

        // table API
        Table resTable = sensorTable
                .groupBy("id")
                .aggregate("avgTemp(temp) as avgTemp")
                .select("id, avgTemp");

        // SQL
        tableEnv.createTemporaryView("sensorTable", sensorTable);
        Table sqlTable = tableEnv.sqlQuery("select id, avgTemp(temp) as avgTemp" +
                " from sensorTable" +
                " group by id");

        // 输出结果
        sensorTable.printSchema();
        tableEnv.toRetractStream(resTable, Row.class).print("resTable");
        tableEnv.toRetractStream(sqlTable, Row.class).print("sqlTable");

        env.execute();
    }

    // TODO: 2022/5/19 实现自定义的 Aggregate Function
    public static class AvgTemp extends AggregateFunction<Double, Tuple2<Double, Integer>> {

        @Override
        public Double getValue(Tuple2<Double, Integer> doubleIntegerTuple2) {
            return doubleIntegerTuple2.f0 / doubleIntegerTuple2.f1;
        }

        @Override
        public Tuple2<Double, Integer> createAccumulator() {
            return new Tuple2<>(0.0, 0);
        }

        // 必须实现accumulate方法
        public void accumulate( Tuple2<Double, Integer> accumulator, Double temp) {
            accumulator.f0 += temp;
            accumulator.f1++;
        }
    }
}
